import gzip
import json


def formal_fid(fid: str):
    return "<http://rdf.freebase.com/ns/" + fid + ">"


def trunc_fid(fid: str):
    fid = fid.replace("<http://rdf.freebase.com/ns/", "").replace(">", "")
    return fid


def judge_fid(fid: str):
    if fid.find("<http://rdf.freebase.com/ns/") == -1:
        return False
    else:
        return True


def douban2fb(pth):
    d2f = {}
    with open(pth, "r", encoding="utf-8") as f:
        for line in f:
            d_id, f_id = line.strip().split()
            d2f[d_id] = f_id
    return d2f


def dict_increase(d: dict, k_list):
    for k in k_list:
        if k in d:
            d[k] += 1
        else:
            d[k] = 1


def dict_decrease(d: dict, k_list):
    for k in k_list:
        d[k] -= 1


def dict_clear(d: dict):  # 删除所有值为0的记录
    for k, v in list(d.items()):
        if v == 0:
            d.pop(k)


def dict_append(d: dict, k, v):
    if k in d:
        d[k].append(v)
    else:
        d[k] = [v]


def shrink_graph(g, g_temp, ent_info, rel_info, req, white_list):
    """
    对子图进行修剪
    g: 子图
    g_temp: 暂存图信息，用于修剪，修剪后加入g
    st_info: 为修剪提供统计信息
    req: 子图修剪的规则, 按[(ent_min,ent_max),(rel_min,rel_max)]的格式记录
    k: 第k轮修剪
    """
    fid_set = set()  # 下一轮中需要查找的结点
    ent_min, ent_max = req[0]
    rel_min, rel_max = req[1]

    delete = True
    while delete:  # 不断循环，直到没有新的元组可以被修剪
        delete = False
        for head, rel, tail in list(g_temp):
            # 判断是否被修剪
            if not (
                (ent_min < ent_info[head] < ent_max)
                and (ent_min < ent_info[tail] < ent_max)
                and (rel_min < rel_info[rel] < rel_max)
            ):  # 应被修剪
                g_temp.remove((head, rel, tail))
                dict_decrease(ent_info, (head, tail))
                dict_decrease(rel_info, (rel,))
                delete = True
        #! len(g_temp)==0 ERROR!
        if len(g_temp) == 0:
            print("ERROR: graph is empty!\n")
            exit(-1)

    for head, rel, tail in g_temp:
        fid_set.update((head, tail))
        dict_append(g, head, (rel, tail))

    return fid_set


def generate_graph(d2f: dict, pth):
    g = {}  # dict: id->list, list[(edge, tail)]
    ent_info = {}  # dict: id->int
    rel_info = {}  # dict: id->int

    n = 2  # 子图跳数
    # 子列表k是第k轮的删减条件，按[(ent_min,ent_max),(rel_min,rel_max)]的格式记录
    req = [[(20, float("inf")), (50, float("inf"))], [(15, 20000), (50, float("inf"))]]
    white_list = set(d2f.values())
    fid_set = set(d2f.values())  # 所有待找 head
    with gzip.open(pth, "rb") as f:
        for i in range(n):
            g_temp = set()  # 暂存图信息，用于修剪，修剪后加入g, 格式为(head, rel, tail)
            #! TODO 目前没想到很好的查找方法，暂且通过遍历一遍 gz 文件实现
            for line in f:
                triplet = line.strip().decode().split("\t")[:3]
                head, rel, tail = triplet
                if judge_fid(head) and judge_fid(rel) and judge_fid(tail):  # 判断fid是否符合要求
                    head, rel, tail = [trunc_fid(fid) for fid in triplet]
                    if head in fid_set:  # 是需要寻找出边的结点
                        g_temp.add((head, rel, tail))
                        dict_increase(ent_info, (head, tail))
                        dict_increase(rel_info, (rel,))
            # dict_save(ent_info, f"./tmp/ent_info_{i}_before.json")
            # dict_save(rel_info, f"./tmp/rel_info_{i}_before.json")

            fid_set = shrink_graph(g, g_temp, ent_info, rel_info, req[i], white_list)  # 修剪子图
            dict_clear(ent_info)
            dict_clear(rel_info)
            dict_save(ent_info, f"./tmp/ent_info_{i}_after.json")
            dict_save(rel_info, f"./tmp/rel_info_{i}_after.json")
            print(f"round {i} over")
            f.seek(0)

    st_info = {"entity": ent_info, "relation": rel_info}
    return g, st_info


def dict_save(d: dict, pth):
    d_js = json.dumps(d, indent=4)
    with open(pth, "w", encoding="utf-8") as f:
        f.write(d_js)


def total_count(g: dict, st_info: dict):
    total = {"entity_num": 0, "relation_num": 0, "triplet_num": 0}
    for l in g.values():
        total["triplet_num"] += len(l)
    total["entity_num"] = len(st_info["entity"])
    total["relation_num"] = len(st_info["relation"])
    return total


if __name__ == "__main__":
    d2f_input_pth = "./data/douban2fb.txt"
    d2f_output_pth = "./res/douban2fb.json"
    db_pth = "./data/freebase_douban.gz"
    graph_pth = "./res/graph_white.json"
    info_pth = "./res/st_info_white.json"

    d2f = douban2fb(d2f_input_pth)
    g, st_info = generate_graph(d2f, db_pth)
    total = total_count(g, st_info)
    st_info.update(total)

    dict_save(d2f, d2f_output_pth)
    dict_save(g, graph_pth)
    dict_save(st_info, info_pth)
